import numpy as np
import scipy.io as sio
import glob
import cv2
from sklearn.preprocessing import MinMaxScaler
import pywt
import matplotlib.pyplot as plt
from scipy.signal import hilbert
from scipy.fft import fft
from scipy.ndimage import grey_dilation, grey_erosion, grey_opening, grey_closing
import scipy.io


'''加载数据'''
# 数据加载
data_paths1 = glob.glob(r"D:\pycharm_project\Bearing_inspection\Be_legend\data\12k Drive End Bearing Fault Data\*.mat")
data_paths2 = glob.glob(r"D:\pycharm_project\Bearing_inspection\Be_legend\data\Normal Baseline Data\*.mat")
data_paths = data_paths1 + data_paths2

# print(data_paths[:5])
ALL_DE_data = []
ALL_DE_label = []
ALL_FE_data = []
ALL_RPM = []
'''数据截取'''
# CWRU数据集中采样频率是12 kHz
fs = 12000
# 去除前5秒的非稳态信号
start_index = 5 * fs

'''数据集分类'''
# 轴承内圈
Bearing_inner_ring = [105, 106, 107, 108, 169, 170, 171, 172, 209, 210, 211, 212, 3001, 3002, 3003, 3004]
# 滚动体
Rolling_element = [118, 119, 120, 121, 185, 186, 187, 188, 222, 223, 224, 225, 3006, 3007, 3008]
# 外圈
Outer_ring = [130, 131, 132, 133, 144, 145, 146, 147, 156, 158, 159, 160, 197, 198, 199, 200, 234, 235, 236, 237, 246,
              247, 248, 249, 258, 259, 260, 261]
# 正常数据
nomal_data = [97, 98, 99, 100]

# 归一化到[-1, 1]
scaler = MinMaxScaler(feature_range=(-1, 1))

'''for data_path in data_paths:
    data_split =data_path.split("")
    # print(data_split[6])
    code = data_split[6].split('.')
    # print(code[0])
    code_num = int(code[0])
    mat_data = sio.loadmat(data_path)
    if code_num == 174:
        DE_time = mat_data['X173_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        FE_time = mat_data['X173_FE_time']  # 风扇端（Fan End, FE）的时域振动信号数据。
        RPM = mat_data['X173RPM']  # 电机转速（Revolutions Per Minute, RPM）
        ALL_DE_data.append(DE_time)
        ALL_FE_data.append(FE_time)
        ALL_RPM.append(RPM)
    elif 100 <= code_num < 999:
        DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        FE_time = mat_data[f'X{code_num}_FE_time']  # 风扇端（Fan End, FE）的时域振动信号数据。
        # 归一化，将数据处理成（-1， 1）之间
        DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
        FE_time = scaler.fit_transform(FE_time.reshape(-1, 1)).flatten()

        RPM = mat_data[f'X{code_num}RPM'].flatten()  # 电机转速（Revolutions Per Minute, RPM）
        ALL_DE_data.append(DE_time[start_index:])
        ALL_FE_data.append(FE_time[start_index:])
        ALL_RPM.append(RPM)
'''

for data_path in data_paths:
    data_split =data_path.split("\\")
    # print(data_split[6])
    code = data_split[6].split('.')
    # print(code[0])
    code_num = int(code[0])
    mat_data = sio.loadmat(data_path)
    if code_num in Bearing_inner_ring:
        DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
        ALL_DE_data.append(DE_time[start_index:])
        ALL_DE_label.append(0)

    elif code_num in Rolling_element:
        DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
        ALL_DE_data.append(DE_time[start_index:])
        ALL_DE_label.append(1)

    elif code_num in Outer_ring:
        DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
        ALL_DE_data.append(DE_time[start_index:])
        ALL_DE_label.append(2)

    elif code_num in nomal_data:
        if code_num == 100:
            DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        else:
            DE_time = mat_data[f'X0{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
        DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
        ALL_DE_data.append(DE_time[start_index:])
        ALL_DE_label.append(3)


data_index = 50
print("DE_data的有效值: " + str(len(ALL_DE_data)))
print("label的有效值: " + str(len(ALL_DE_label)))
# print("ALL_FE_data的有效值: " + str(len(ALL_FE_data)))
print(ALL_DE_data[data_index])
print(len(ALL_DE_data[data_index]))

def normalize_to_uint8(M):
    # M in [-1,1] → [0,255]
    M2 = ((M + 1) * 127.5).clip(0, 255)
    return M2.astype(np.uint8)

def GADF(signal):
    diff = signal[:, None] - signal[None, :]
    M = np.sin(diff)  # GADF 差场矩阵，shape = (h, h)
    M = normalize_to_uint8(M)
    return M


def generate_cwt(signal, scales=np.arange(1, 65)):
    # 使用cmor小波（论文推荐）
    coefficients, _ = pywt.cwt(signal, scales, 'cmor1.5-1.0')
    # 取模值并归一化
    cwt_image = np.abs(coefficients)
    cwt_image = (cwt_image - np.min(cwt_image)) / (np.max(cwt_image) - np.min(cwt_image))

    # 输出为(len(scales)×len(signal))的二维矩阵
    return cwt_image

'''数据分段'''
# 分段参数设置
phi_segments_line = ALL_DE_data[data_index]
phi_segments_line = np.arccos(phi_segments_line)
window_size = 64      # 窗口长度
num_windows = phi_segments_line.shape[0] // window_size  # 向下取整

# 存放所有 GADF 图像（也可以直接保存成文件，不一定要一次性存到内存）
gadf_images = []

for num_window in range(num_windows):
    start = num_window * window_size
    end = start + window_size
    '''
    # 格拉姆角场方法（GADF）
    M = GADF(phi_segments_line, start, end)
    '''
    # 连续小波变换方法（CWT）
    M = generate_cwt(phi_segments_line[start:end])
    gadf_images.append(M)

print(f"生成了 {len(gadf_images)} 张 {window_size}×{window_size} 的 GADF 图")
print(gadf_images[1].shape)


# --- 初始化索引 ---
index = 0
max_index = len(gadf_images) - 1
# t = np.linspace(0, 1, 1024)
# scales = np.arange(1, 100)
# # 显示结果
# plt.figure(figsize=(8, 4))
# plt.imshow(gadf_images[1], aspect='auto', cmap='viridis',
#            extent=[t[0], t[-1], scales[-1], scales[0]])
# plt.colorbar(label='振幅')
# plt.title("Correct CWT Image (论文结果)")
# plt.show()


# ---- 主循环 ----
while True:
    # 取出当前 GADF 矩阵
    img = gadf_images[index]           # shape (h, h) 或 (h, h, 1)
    img_large = cv2.resize(img, (256, 256), interpolation=cv2.INTER_LINEAR)
    # img = normalize_to_uint8(M)
    cv2.imshow('GADF Viewer', img_large)
    key = cv2.waitKey(1)
    if key != -1:                       # 255 表示无按键
        print("Pressed key code:", key)  # 调试：看看捕获到的值
        print(img.shape)

    if key == 111:
        # 按 'o' 退出
        break

    elif key == 113:
        # 按 'q' 浏览上一张，环绕
        index = (index - 1) % (max_index + 1)

    elif key == 101:
        # 按 'w' 浏览下一张，环绕
        index = (index + 1) % (max_index + 1)

# 循环结束后关闭所有窗口
cv2.destroyAllWindows()


'''############################# 基础信息 ##############################'''
'''matplotlib解决中文字体问题'''
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体（如黑体）
plt.rcParams['axes.unicode_minus'] = False   # 解决负号显示异常

def function_EMTF(y, g_size):
    """
    多尺度形态学增强函数（等价于 MATLAB 的 function_EMTF）
    y: 原始信号 (1D array)
    g_size: 结构元素的尺度 (int)
    return: 增强后的信号
    """

    # 创建结构元素
    structure = np.ones(g_size) # 卷积核，全 1 数组，类似于滑动窗口

    # 基本形态学操作
    fd = grey_dilation(y, structure=structure)  # 膨胀操作，凸显峰值
    fe = grey_erosion(y, structure=structure)   # 蚀操作，凸显谷值
    fo = grey_opening(y, structure=structure)   # 开运算 = 腐蚀后膨胀（滤除小波峰
    fc = grey_closing(y, structure=structure)   # 闭运算 = 膨胀后腐蚀（滤除小波谷）

    fco = grey_opening(fc, structure=structure)
    fcd = grey_dilation(fc, structure=structure)

    foc = grey_closing(fo, structure=structure)
    foe = grey_erosion(fo, structure=structure)

    fdc = grey_closing(fd, structure=structure)

    feo = grey_opening(fe, structure=structure)

    # 提取信号中的调制特征和结构性变化特征，用于后续故障识别、分类或者可视化
    C = fcd - feo
    D = fdc - foe
    E = (C + D) / 2

    fAco_oc = (fco + foc) / 2
    fCMFH = y - fAco_oc

    DIF = fc - fo
    fcd_oe = (fcd + foe) / 2
    Fcd_oe = y - fcd_oe
    MG = fd - fe
    WT = y - (fd + fe) / 2

    AA = fco - fe
    BB = foc - fd
    AVG = (fo + fc) / 2
    AVGH = y - AVG

    fcoe = grey_erosion(grey_opening(fc, structure=structure), structure=structure)
    focd = grey_dilation(grey_closing(fo, structure=structure), structure=structure)
    fdeo = grey_opening(grey_erosion(fd, structure=structure), structure=structure)
    fedc = grey_closing(grey_dilation(fe, structure=structure), structure=structure)

    fcde = grey_erosion(grey_dilation(fc, structure=structure), structure=structure)
    foed = grey_dilation(grey_erosion(fo, structure=structure), structure=structure)
    fdce = grey_erosion(grey_closing(fd, structure=structure), structure=structure)
    feod = grey_dilation(grey_opening(fe, structure=structure), structure=structure)

    A = fcoe - focd
    B = fdeo - fedc

    W = fcde - foed
    fcde_oed = (fcde + foed) / 2
    w = y - fcde_oed

    # 最终输出（乘法增强残差）
    EMTF = 2 * WT * w

    return EMTF


def main():
    # === 1. 载入数据 ===
    mat = sio.loadmat(r"D:\pycharm_project\Bearing_inspection\Be_legend\data\12k Drive End Bearing Fault Data\51.mat")  # 包含变量 y
    y = mat['X051_DE_time'].squeeze()
    y = y[:256]
    fs = 12000            # 采样频率
    N = 256              # 取样点数
    '''用于信号的频谱分析准备工作'''
    t = np.arange(N) / fs  # 时间向量，对应每一点采集的时间刻度
    L = len(t)
    NFFT = 2 ** int(np.ceil(np.log2(L)))  # 最近的 2 的幂次，确保FFT的长度是2的幂次方（比如4096）
    f = fs / 2 * np.linspace(0, 1, NFFT // 2 + 1)  # 单边频率向量，每个频率点的实际频率（Hz），只处理前半段，因为前半段和后半段共轭对称

    scale_range = range(1, 300)         # 一共 199 个尺度
    zatL = np.zeros((len(scale_range), N))  # [199, 4096]，存储每个尺度处理后的信号
    baoluoatL = np.zeros((len(scale_range), len(f)))    # 储每个尺度下信号的包络谱，即信号的特征

    # === 2. 多尺度形态学增强 + 包络谱 ===
    for i, s in enumerate(scale_range):
        g_size = s + 1
        zatL[i, :] = function_EMTF(y, g_size)

    for i in range(len(scale_range)):
        analytic_signal = hilbert(zatL[i, :])
        envelope = np.abs(analytic_signal)
        envelope_fft = fft(envelope, n=NFFT) / L
        baoluoatL[i, :] = 2 * np.abs(envelope_fft[:NFFT // 2 + 1])

    # 去掉直流分量
    baoluoatL[:, 0] = 0

    # === 3. 提取并可视化频率-尺度图谱 ===
    B = list(range(len(scale_range)))  # 修正索引问题
    A = baoluoatL[B, :]
    C = A[:, :87]
    F = f[:87]

    # 假设 C 是 shape = [199, 87] 的二维数组
    C_normalized = (C - np.min(C)) / (np.max(C) - np.min(C))  # 归一化到 [0,1]
    image_array = (C_normalized * 255).astype(np.uint8)  # 转为 uint8 图像格式
    image_rgb = np.stack([image_array] * 3, axis=-1)  # shape: (199, 87, 3)
    print(image_rgb.shape)
    cv2.imshow("ss", image_rgb)
    key = cv2.waitKey(1)


if __name__ == "__main__":
    main()


